import numpy as np
import matplotlib.pyplot as plt
# 生成数据
stock_cnt = 200
view_days = 504
stock_day_change = np.random.standard_normal((stock_cnt, view_days))
print(stock_day_change.shape)
# 0:2第一，第二支股票，0:5头五个交易日的涨跌幅数据
print(stock_day_change[0:2, :5])
print(stock_day_change[-2:, -5:])

# 计算
import scipy.stats as scs
stock_mean = stock_day_change[0].mean()
stock_std = stock_day_change[0].std()
print('{:.3f}'.format(stock_mean))
print('{:.3f}'.format(stock_std))

plt.hist(stock_day_change[0], bins=505, density=True) # 这个参数指定bin(箱子)的个数,也就是总共有几条条状图
fit_linspace = np.linspace(stock_day_change[0].min(), stock_day_change[0].max())
pdf = scs.norm(stock_mean, stock_std).pdf(fit_linspace)
plt.plot(fit_linspace, pdf, lw=2, c='r')
# plt.show()
# np.save('./gen/stock_day_change', stock_day_change)
stock_day_change = np.load('./gen/stock_day_change.npy')
stock_day_change.shape
# 策略

keep_days = 50
stock_day_change_test = stock_day_change[:stock_cnt, 0:view_days - keep_days]
# 从200 支股票，切片取出0 ~ 504-50 数据

# 打印出前454跌幅最大的三支，总跌幅通过np.sum计算，np.sort对结果排序
# 因为从小到大排序，取前三就代表跌幅最大
print(np.sort(np.sum(stock_day_change_test, axis=1))[:3])

# test = np.sort(np.sum(stock_day_change_test, axis=1))
# print(test.shape)
# 使用np.argsort针对股票跌幅进行排序，返回序号，即符合买入条件的股票序号
stock_lower_array = np.argsort(np.sum(stock_day_change_test, axis=1))[:3]
# 输出符合买入条件的股票序号
print(stock_lower_array)



def show_bug_lower(stock_ind):
    _, axs = plt.subplots(nrows=1, ncols=3, figsize=(16, 5))

    axs[0].plot(np.arange(0, view_days - keep_days),
                stock_day_change_test[stock_ind].cumsum())

    # if stock_ind == 142:
    print("hhhd")
    # axs[2].plot(np.arange(0, view_days), stock_day_change[stock_ind].cumsum())
    print('{:.3f}'.format(stock_day_change[stock_ind].mean()))
    # 142支股票均值为-0.139，504天将会跌至70元，如果依照前454天的趋势买入股票，由于前454天只下跌了50元，
    # 后面还会继续下跌，所以这只股票该策略会亏损

    # 变化率的累加看出趋势
    cs_buy = stock_day_change[stock_ind][view_days - keep_days:view_days].cumsum()
    axs[1].plot(np.arange(view_days - keep_days, view_days), cs_buy)
    plt.show()
    return cs_buy[-1]

profit = 0
for stock_ind in stock_lower_array:
    profit += show_bug_lower(stock_ind)

print('{}, 盈亏{:.2f}%'.format(stock_lower_array, profit))

